# Licensed under the MIT license.

"""
rStar 命令行参数配置模块

本模块定义了 rStar 项目的完整参数配置系统，包括：
1. 基础参数解析器配置 - 模型、数据集、API等核心设置
2. 参数后处理逻辑 - 输出目录创建、GPU检查、模型兼容性调整
3. 参数保存机制 - 实验配置的持久化存储

主要功能：
- 统一管理所有命令行参数的定义和默认值
- 自动处理不同模型(特别是Qwen3)的兼容性问题
- 创建结构化的输出目录便于结果管理
- 提供GPU环境检测和验证

作者: rStar团队
创建时间: 2024
"""

import os, json, torch, math
from argparse import ArgumentParser
from datetime import datetime


def get_parser():
    """
    创建并配置基础命令行参数解析器
    
    构建 rStar 项目的核心参数配置，涵盖以下主要类别：
    1. 实验配置 - 实验标识、随机种子、详细输出控制
    2. 模型配置 - 模型路径、API选择、并行设置、推理参数
    3. 数据配置 - 数据集选择、文件路径、处理范围
    4. 输出配置 - 结果保存路径、日志记录设置
    5. 特殊功能 - WandB集成、特定模型优化等
    
    Returns:
        ArgumentParser: 配置完成的参数解析器，包含所有基础参数定义
        
    参数分类详解：
        - 必需参数: model_ckpt, dataset_name (用户必须指定)
        - 模型参数: api, tensor_parallel_size, max_tokens等 (控制推理行为)
        - 数据参数: data_root, test_json_filename等 (控制数据加载)
        - 系统参数: seed, verbose, model_parallel等 (控制系统行为)
    """
    parser = ArgumentParser()

    # ================================ 实验标识与控制参数 ================================
    parser.add_argument(
        "--note", 
        type=str, 
        default="debug",
        help="实验备注标识。用于区分不同实验配置，会被添加到输出目录名中。"
             "建议使用描述性名称如 'baseline_test', 'ablation_a1_disabled' 等"
    )

    # ================================ API与模型配置 ================================
    allowed_apis = ["together", "huggingface", "llama", "vllm", "debug", "gpt3.5-turbo"]
    parser.add_argument(
        "--api", 
        type=str, 
        choices=allowed_apis, 
        default="vllm", 
        help=f"推理API类型选择: {allowed_apis}。"
             "vllm: 高性能推理引擎(推荐), huggingface: 原生transformers, "
             "gpt3.5-turbo: OpenAI云端API, debug: 调试模式"
    )

    # ================================ 基础系统参数 ================================
    parser.add_argument(
        "--seed", 
        type=int, 
        default=42,
        help="随机种子。确保实验结果的可重现性。"
             "相同种子在相同环境下应产生一致的结果"
    )
    
    parser.add_argument(
        "--verbose", 
        action="store_true",
        help="启用详细输出模式。显示MCTS搜索过程、推理步骤、资源使用等调试信息。"
             "建议在开发和调试阶段启用，生产环境可关闭以提高性能"
    )

    # ================================ WandB实验跟踪配置 ================================
    parser.add_argument(
        "--wandb_mode", 
        type=str, 
        default="disabled", 
        choices=["disabled", "online"],
        help="Weights & Biases实验跟踪模式。"
             "online: 启用在线同步和可视化, disabled: 禁用WandB功能"
    )

    # ================================ 模型加载与推理配置 ================================
    parser.add_argument(
        "--model_ckpt", 
        required=True,
        help="模型检查点路径(必需参数)。支持以下格式："
             "本地路径: /path/to/model, HuggingFace模型: microsoft/DialoGPT-medium, "
             "OpenAI模型: gpt-3.5-turbo。路径必须存在且可访问"
    )

    parser.add_argument(
        "--model_parallel", 
        action="store_true",
        help="启用模型并行加载。适用于大模型需要多GPU分布式加载的场景。"
             "自动检测可用GPU数量并配置分布式环境"
    )
    
    parser.add_argument(
        "--half_precision", 
        action="store_true",
        help="启用半精度(FP16)推理。显著降低显存占用和推理延迟，略微影响精度。"
             "对于大模型或GPU内存受限环境强烈推荐启用"
    )

    # ================================ 推理生成参数 ================================
    parser.add_argument(
        "--max_tokens", 
        type=int, 
        default=1024, 
        help="单次生成的最大token数量。控制模型输出长度上限。"
             "过小可能截断推理过程，过大可能产生冗余内容。推荐范围: 512-2048"
    )
    
    parser.add_argument(
        "--temperature", 
        type=float, 
        default=0.8, 
        help="生成温度参数。控制输出的随机性和多样性。"
             "0.0: 贪心选择(确定性), 1.0: 完全随机, 0.7-0.9: 平衡创造性与一致性"
    )
    
    parser.add_argument(
        "--top_k", 
        type=int, 
        default=40, 
        help="Top-K采样参数。限制候选token为概率最高的K个。"
             "较小值提高确定性，较大值增加多样性。推荐范围: 20-50"
    )
    
    parser.add_argument(
        "--top_p", 
        type=float, 
        default=0.95, 
        help="Top-P(nucleus)采样参数。基于累积概率选择候选token。"
             "0.9-0.95适合大多数任务，更高值增加多样性但可能降低质量"
    )
    
    parser.add_argument(
        "--num_beams", 
        type=int, 
        default=1, 
        help="束搜索(beam search)的束宽。1表示贪心搜索，>1启用束搜索。"
             "增加束宽可能提高质量但显著增加计算成本"
    )

    # ================================ 系统资源与并行配置 ================================ 
    parser.add_argument(
        "--max_num_worker", 
        type=int, 
        default=4, 
        help="数据加载器的最大工作线程数。影响数据预处理和加载速度。"
             "过多线程可能导致内存压力，建议根据CPU核心数调整"
    )
    
    parser.add_argument(
        "--test_batch_size", 
        type=int, 
        default=1,
        help="测试批次大小。当前MCTS推理主要支持单样本处理。"
             "保持默认值1，未来版本可能支持批量处理"
    )
    
    parser.add_argument(
        "--tensor_parallel_size", 
        type=int, 
        default=1,
        help="张量并行度。将模型分割到多个GPU上进行并行推理。"
             "必须是GPU总数的因子，如8卡可设置为1,2,4,8。大模型推荐使用"
    )
    
    parser.add_argument(
        "--max_num_seqs", 
        type=int, 
        default=256, 
        help="vLLM引擎的最大并发序列数。控制批量推理的并发度。"
             "较大值提高吞吐量但增加显存占用。根据GPU内存和模型大小调整"
    )

    # ================================ 提示词与数据配置 ================================
    parser.add_argument(
        "--prompts_root", 
        default="prompts",
        help="提示词模板根目录。包含各数据集的Few-shot示例和推理模板。"
             "目录结构: prompts/{dataset_name}/{prompt_type}/"
    )

    parser.add_argument(
        "--data_root", 
        default="data",
        help="数据集根目录。包含所有支持的数学推理数据集。"
             "目录结构: data/{dataset_name}/{split}.json"
    )
    
    allowed_dataset_names = ["MATH", "GSM8K", "GSM8KHARD", "STG", "SVAMP", "MULTIARITH"]
    parser.add_argument(
        "--dataset_name",
        required=True,
        choices=allowed_dataset_names,
        help=f"目标数据集名称(必需参数): {allowed_dataset_names}。"
             "GSM8K: 小学数学, MATH: 高中竞赛数学, MULTIARITH: 多步算术, "
             "GSM8KHARD: GSM8K困难版, STG/SVAMP: 数学应用题变体"
    )
    
    parser.add_argument(
        "--test_json_filename", 
        type=str, 
        default="test_all",
        help="测试数据文件名(不含.json扩展名)。"
             "完整路径: {data_root}/{dataset_name}/{test_json_filename}.json"
    )
    
    parser.add_argument(
        "--start_idx", 
        type=int, 
        default=0, 
        help="测试问题的起始索引(包含)。支持分片处理大数据集。"
             "与end_idx配合使用可以并行处理不同数据片段"
    )
    
    parser.add_argument(
        "--end_idx", 
        type=int, 
        default=math.inf, 
        help="测试问题的结束索引(不包含)。默认math.inf表示处理到数据集末尾。"
             "与start_idx配合实现数据集分片处理"
    )

    # ================================ 输出配置 ================================
    parser.add_argument(
        "--outputs_dir", 
        type=str, 
        default="./run_outputs",
        help="输出目录。存储生成的解答、中间结果和统计信息。"
             "自动创建时间戳子目录以区分不同运行"
    )

    # ================================ 特殊模型优化参数 ================================
    parser.add_argument(
        "--qwen3_disable_thinking", 
        action="store_true", 
        help="禁用Qwen3模型的thinking模式。解决Qwen3特有的输出格式问题。"
             "检测到Qwen3模型时会自动启用此选项"
    )
    
    parser.add_argument(
        "--enable_potential_score", 
        action="store_true", 
        help="启用基于潜在分数的答案选择机制。通过预测答案质量提高选择准确性。"
             "增加计算开销但可能提升最终答案质量"
    )

    return parser


def post_process_args(args):
    """
    参数后处理：创建输出目录、环境检查和模型兼容性调整
    
    执行以下关键后处理步骤：
    1. 输出目录管理 - 创建带时间戳的结构化输出目录
    2. GPU环境检测 - 验证可用GPU并记录设备信息  
    3. 模型兼容性 - 针对特定模型(如Qwen3)应用兼容性设置
    4. 参数验证 - 检查参数组合的合理性和可行性
    
    Args:
        args (Namespace): 解析后的命令行参数对象，包含用户指定的所有配置
        
    Returns:
        Namespace: 后处理完成的参数对象，包含：
            - 完整的输出路径配置 (run_outputs_dir, answer_sheets_dir等)
            - GPU设备信息 (cuda_0, cuda_1等)  
            - 兼容性调整后的模型参数
            
    主要处理逻辑：
        - 目录创建：基于时间戳和实验标识创建唯一输出目录
        - GPU检测：识别可用GPU型号并提供优化建议
        - Qwen3优化：自动应用Qwen3模型的最佳实践配置
        - 环境变量：设置模型推理所需的系统环境变量
    """
    # ================================ 输出目录结构配置 ================================
    # 构建带时间戳和实验标识的输出目录名
    suffix = "---[" + args.note + "]" if args.note is not None else ""
    model_name = args.model_ckpt.split("/")[-1]  # 提取模型名称作为目录标识
    
    # 创建主输出目录：{outputs_dir}/{dataset}/{model}/{timestamp}---[note]
    args.run_outputs_dir = os.path.join(
        args.outputs_dir,
        args.dataset_name,           # 数据集级别分类
        model_name,                  # 模型级别分类  
        f"{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}" + suffix  # 时间戳+标识
    )
    os.makedirs(args.run_outputs_dir, exist_ok=True)

    # 创建答案保存子目录：存储每个问题的详细解答结果
    args.answer_sheets_dir = os.path.join(args.run_outputs_dir, "answer_sheets")
    os.makedirs(args.answer_sheets_dir, exist_ok=True)

    # ================================ Qwen3模型兼容性优化 ================================
    """
    Qwen3模型需要特殊的配置优化以确保稳定运行：
    1. 禁用thinking模式避免输出格式问题
    2. 降低并发数量减少内存压力
    3. 启用半精度推理提高效率
    4. 设置专用环境变量优化性能
    """
    if "Qwen3" in args.model_ckpt:
        print("检测到Qwen3模型，正在应用兼容性设置...")
        
        # 自动启用thinking模式禁用，解决输出格式不兼容问题
        if not args.qwen3_disable_thinking:
            print("自动启用 --qwen3_disable_thinking")
            args.qwen3_disable_thinking = True
        
        # 降低最大并发序列数，避免Qwen3的内存管理问题
        if args.max_num_seqs > 16:
            print(f"为Qwen3模型调整 max_num_seqs: {args.max_num_seqs} -> 16")
            args.max_num_seqs = 16
        
        # 强制启用半精度推理，提高Qwen3的内存效率
        if not args.half_precision:
            print("为Qwen3模型启用半精度")
            args.half_precision = True
        
        # 设置Qwen3专用环境变量，优化推理性能和稳定性
        os.environ["VLLM_TORCH_COMPILE_LEVEL"] = "0"           # 禁用torch编译
        os.environ["VLLM_DISABLE_CUSTOM_ALL_REDUCE"] = "1"     # 禁用自定义all-reduce
        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"  # 启用内存段扩展
        print("已设置Qwen3兼容性环境变量")

    # ================================ GPU环境检测与验证 ================================
    """
    检测并记录可用GPU信息，为用户提供硬件配置建议：
    1. 统计可用GPU数量和型号
    2. 针对特定GPU型号提供优化建议
    3. 存储GPU信息供后续配置参考
    """
    num_gpus = torch.cuda.device_count()
    cuda_devices = [torch.cuda.get_device_name(i) for i in range(num_gpus)]
    
    # 确保至少有一块可用GPU
    assert len(cuda_devices) > 0, "错误：未检测到可用GPU，rStar需要GPU支持"
    
    # 记录各GPU设备信息（最多支持4卡）
    args.cuda_0 = cuda_devices[0]
    args.cuda_1 = cuda_devices[1] if len(cuda_devices) > 1 else None
    args.cuda_2 = cuda_devices[2] if len(cuda_devices) > 2 else None
    args.cuda_3 = cuda_devices[3] if len(cuda_devices) > 3 else None

    # ================================ GPU优化建议 ================================
    # 针对A100-SXM4-40GB提供优化建议：建议启用半精度以充分利用硬件
    if len(cuda_devices) == 1:
        if args.cuda_0 == "NVIDIA A100-SXM4-40GB" and not args.half_precision:
            print("建议：检测到A100-SXM4-40GB，推荐启用 --half_precision 以提高性能")

    return args


def save_args(args):
    """
    保存参数配置到JSON文件
    
    将完整的参数配置持久化存储，便于：
    1. 实验重现 - 记录完整的配置信息
    2. 结果分析 - 关联配置与性能表现  
    3. 调试追踪 - 快速定位配置问题
    4. 批量实验 - 比较不同配置的效果
    
    Args:
        args (Namespace): 后处理完成的参数对象
        
    生成文件：
        {run_outputs_dir}/args.json: 包含所有参数的JSON配置文件
        
    文件内容：
        完整的键值对映射，包括默认值、用户设置和后处理结果
        便于后续程序读取和人工检查
    """
    # 将Namespace对象转换为字典并保存为JSON格式
    args_dict = vars(args)
    args_file_path = os.path.join(args.run_outputs_dir, "args.json")
    
    with open(args_file_path, "w", encoding='utf-8') as f:
        json.dump(args_dict, f, indent=4, ensure_ascii=False)
    
    print(f"参数配置已保存到: {args_file_path}")
